# -*- coding: utf-8 -*-
"""
Created on Wed Apr 13 10:58:42 2022

@author: 11325
"""

import torch
import numpy as np
import pandas as pd
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import dataloader 
import Evaluation 
import TransModel
import os


def train(model,dataload,loss_fuc,train_data,test_data,learning_rate,num_epoch,batch_size):
    optimizer=torch.optim.SGD(model.parameters(), lr=learning_rate)
    loss_sum=torch.cuda.FloatTensor([0.0])
    model.model_normalize()
    for e in range(num_epoch): 
        loss_sum=0
        print("Epoch {}".format(e))
        for correct_tuple, corrupt_tuple in dataload.data_iter(train_data,batch_size):
            correct_score,corrput_score=model(correct_tuple, corrupt_tuple)           
            loss=loss_fuc(correct_score,corrput_score,1)
            model.zero_grad()
            loss.backward()
            optimizer.step()
            loss_sum+=loss.data.item()
        print("Average loss:{}".format((loss_sum/len(train_data))))
    print('Training complete and save model')
    torch.save(model,dataloader.model_save_path)
        


if __name__=='__main__' : 
    print("Loading train and test triplets......")
    train_path=dataloader.FB15K_BASE_PATH_train_PATH
    test_path=dataloader.FB15K_BASE_PATH_1_1_PATH
    train_enity,train_realtion,train_triple=dataloader.get_pairs(train_path)
    _,_,test_pair=dataloader.get_pairs(test_path)
    test_triple=dataloader.load_without_dic(test_pair)
    print('\n Load complete {} train traiples {} test_triples'.format(len(train_triple), len(test_triple)))
    
    print('Intialize model and dataloader')
    dataLoad=dataloader.DataLoad(train_enity, train_enity)
    model=TransModel.TransE(len(train_enity), len(train_realtion), margin=1, embedding_dim=200).cuda()
    Loss_function=Evaluation.Hingleloss().cuda()
    torch.cuda.empty_cache()
    
    print('Start training...')
    train(model, dataLoad, Loss_function, train_triple, test_triple, learning_rate=0.1, num_epoch=10, batch_size=1000)
    
    print('Start testing...')
    test_eve=Evaluation.Evaluation(model, test_triple)
    
    hit10_enlity,mean_rank_enlity=test_eve.test()
    print('Hit10: {} Mean_rank: {}'.format(hit10_enlity,mean_rank_enlity))
    
    